'''
Script to remove the imaginary component from Bruker NMR data

Application -> Boost S/N in phase-corrected SHARPER spectra

    @author: Patrick Boaler
    @email:  Patrick.Boaler@ed.ac.uk

Works with TopSpin 3.5, 4.0 and 4.1. 
Only NMR experiment directories (expnos)can be in the starting directory of this script
'''

import argparse
import os
import shutil
import nmrglue as ng

parser = argparse.ArgumentParser(prog = "zim.py", 
                                 description = 'Script to remove the imaginary component from Bruker NMR data to boost S/N in phase-corrected SHARPER spectra',
                                 epilog = "If no --expno is specified, the script will search the directory for all experiments whose pulse-programme contains \"sharp\" and complete the operation for all of them, depositing the output in a new directory, appended with \"-zim\" ")

parser.add_argument('filepath')
parser.add_argument('-e', '--expno', type=int, help="specify an experiment number to remove imaginary component on a particular experiment (must be an integer)")
args = parser.parse_args()
path = args.filepath

def read_data2(dir):
    try:
        dic,data = ng.bruker.read(dir)
    except ZeroDivisionError:
        try:
            dic,data = ng.bruker.read(dir,'ser')
        except (FileNotFoundError,ZeroDivisionError):
            print('Error on ' + dir)
            return None,None
    return dic,data

if args.expno:
    inpath = os.path.join(path,str(args.expno))
    outpath = os.path.join(path + '-zim',str(args.expno))
    
    if 'pulseprogram' not in os.listdir(inpath):
            shutil.copy(os.path.join(inpath,'pulseprogram.precomp'),os.path.join(inpath,'pulseprogram'))
    
    dic,data = ng.bruker.read(inpath)     # reads bruker data
    newdata = ng.proc_base.di(data) + 0j  # sets imaginary component to 0j for all points in FID
        
    try:
        shutil.copytree(inpath,outpath)
    except FileExistsError:
        shutil.rmtree(outpath)
        shutil.copytree(inpath,outpath)
    ng.fileio.bruker.write(outpath,dic,newdata,overwrite=True)

else:
    dirs = []
    for subdir in os.listdir(path):
        if 'pulseprogram' not in os.listdir(os.path.join(path,subdir)):
            shutil.copy(os.path.join(path,subdir,'pulseprogram.precomp'),os.path.join(path,subdir,'pulseprogram'))
            
        dic,data = read_data2(os.path.join(path,subdir))
        try:
            if 'sharp' in dic['acqus']['PULPROG'].lower():
                dirs.append(subdir)
        except TypeError:
            pass

    for subdir in dirs:
        inpath = os.path.join(path,subdir)
        outpath = os.path.join(path + '-zim',subdir)
        dic,data = read_data2(inpath)
        newdata = ng.proc_base.di(data) + 0j
        try:
            shutil.copytree(inpath,outpath)
            ng.fileio.bruker.write(outpath,dic,newdata,overwrite=True)
        except FileExistsError:
            shutil.rmtree(outpath)
            shutil.copytree(inpath,outpath)
        ng.fileio.bruker.write(outpath,dic,newdata,overwrite=True)

